import torch
import torchvision
import torch.nn as nn
from pretrainedmodels import models
from inceptionv4 import InceptionV4

def model_resnext50_32x4d(Num_Class,pretrained = False,freeze_weight = False):
    ## freeze weight 控制是否冻结参数。
    model = torchvision.models.resnext50_32x4d(num_classes = Num_Class)
    if pretrained:
        pretrained_state_dict = torch.load('../pre_trained_models/resnext50_32x4d.pth')
        pretrained_state_dict = {k:v for k,v in pretrained_state_dict.items() if 'fc' not in k}
        model_state_dict = model.state_dict()
        model_state_dict.update(pretrained_state_dict)
        model.load_state_dict(model_state_dict)
    ## 筛选参数进行冻结
    if freeze_weight:
        for k,v in model.named_parameters():
            if 'fc' not in k:
                v.requires_grad = False
    return model.cuda()


def model_InceptionV4(num_class,pretrained = False):
    model = InceptionV4(num_classes=num_class)
    if pretrained:
        pre_trained_state_dict = torch.load('../pre_trained_models/inceptionv4.pth')
        model_state_dict = model.state_dict()
        pre_trained_state_dict = {k:v for k,v in pre_trained_state_dict.items()
                                  if 'last_linear' not in k}
        model_state_dict.update(pre_trained_state_dict)
        model.load_state_dict(model_state_dict)
    return model.cuda()



def model_Xception(num_class,pretrained = False):
    model = models.xception(num_classes=num_class,pretrained=False)
    if pretrained:
        pre_trained_state_dict = torch.load('../pre_trained_models/xception.pth')
        model_state_dict = model.state_dict()
        pre_trained_state_dict = {k:v for k,v in pre_trained_state_dict.items()
                                  if 'fc' not in k}
        model_state_dict.update(pre_trained_state_dict)
        model.load_state_dict(model_state_dict)
    return model.cuda()

def model_res18(num_class,pretrained = False):
    model = models.resnet18(num_classes=num_class,pretrained=None)
    if pretrained:
        pre_trained_state_dict = torch.load('../pre_trained_models/resnet18.pth')
        model_state_dict = model.state_dict()
        pre_trained_state_dict = {k:v for k,v in pre_trained_state_dict.items()
                                  if 'fc' not in k}
        model_state_dict.update(pre_trained_state_dict)
        model.load_state_dict(model_state_dict)
    return model.cuda()

if __name__ == '__main__':
    model = model_res18(num_class=3,pretrained=False)
    input = torch.randn(1,3,512,512).cuda()
    out  = model(input)
    print(out.shape)
